"""
实时音频流传输模块
"""

import asyncio
import io
import time
from typing import Dict, List, Any, Optional, Callable, AsyncGenerator
from dataclasses import dataclass, field
from enum import Enum
import numpy as np
import soundfile as sf
import logging
from concurrent.futures import ThreadPoolExecutor


class StreamEventType(Enum):
    """流事件类型"""
    BUFFER_LOW = "buffer_low"
    BUFFER_FULL = "buffer_full"
    BUFFER_EMPTY = "buffer_empty"
    QUALITY_ADJUSTED = "quality_adjusted"
    RECONNECTED = "reconnected"
    ERROR = "error"


@dataclass
class StreamEvent:
    """流事件"""
    event_type: StreamEventType
    timestamp: float
    data: Dict[str, Any] = field(default_factory=dict)


@dataclass
class StreamQuality:
    """流质量设置"""
    sample_rate: int = 22050
    channels: int = 1
    bit_depth: int = 16
    format: str = "wav"
    compression_level: int = 5
    target_bitrate: Optional[int] = None


@dataclass
class AdaptiveConfig:
    """自适应配置"""
    enable_adaptive: bool = True
    min_buffer_size: int = 3
    max_buffer_size: int = 20
    quality_levels: List[StreamQuality] = field(default_factory=lambda: [
        StreamQuality(16000, 1, 16, "wav", 3, 64),
        StreamQuality(22050, 1, 16, "wav", 5, 128),
        StreamQuality(44100, 2, 16, "wav", 7, 256),
    ])
    network_threshold: float = 0.8  # 网络质量阈值


class RealtimeAudioStreamer:
    """实时音频流传输器"""
    
    def __init__(self, config: AdaptiveConfig):
        self.config = config
        self._is_streaming = False
        self._is_paused = False
        self._current_quality = config.quality_levels[1]  # 默认中等质量
        
        # 缓冲区管理
        self._buffer: List[bytes] = []
        self._max_buffer_size = 10
        self._buffer_lock = asyncio.Lock()
        
        # 事件系统
        self._event_handlers: Dict[StreamEventType, List[Callable]] = {}
        self._event_queue = asyncio.Queue()
        
        # 统计信息
        self._stats = {
            "bytes_sent": 0,
            "chunks_sent": 0,
            "buffer_underruns": 0,
            "quality_changes": 0,
            "reconnections": 0
        }
        
        # 线程池
        self._executor = ThreadPoolExecutor(max_workers=2)
        
        # 自适应控制
        self._network_quality = 1.0
        self._buffer_health = 0.5
        
    async def start_streaming(
        self,
        audio_source: AsyncGenerator[bytes, None],
        client_callback: Callable[[bytes], Any]
    ) -> bool:
        """开始实时流传输"""
        try:
            self._is_streaming = True
            self._is_paused = False
            
            # 启动事件处理器
            event_task = asyncio.create_task(self._process_events())
            
            # 启动自适应控制
            adaptive_task = asyncio.create_task(self._adaptive_control())
            
            # 启动流传输
            stream_task = asyncio.create_task(
                self._stream_loop(audio_source, client_callback)
            )
            
            # 等待任一任务完成
            done, pending = await asyncio.wait(
                [event_task, adaptive_task, stream_task],
                return_when=asyncio.FIRST_COMPLETED
            )
            
            # 取消未完成的任务
            for task in pending:
                task.cancel()
                try:
                    await task
                except asyncio.CancelledError:
                    pass
            
            return True
            
        except Exception as e:
            await self._emit_event(StreamEventType.ERROR, {"error": str(e)})
            return False
        finally:
            self._is_streaming = False
    
    async def _stream_loop(
        self,
        audio_source: AsyncGenerator[bytes, None],
        client_callback: Callable[[bytes], Any]
    ):
        """流传输循环"""
        try:
            async for audio_chunk in audio_source:
                if not self._is_streaming:
                    break
                
                if self._is_paused:
                    await asyncio.sleep(0.1)
                    continue
                
                # 根据当前质量处理音频
                processed_chunk = await self._process_audio_chunk(audio_chunk)
                
                if processed_chunk:
                    # 发送到客户端
                    await self._send_to_client(processed_chunk, client_callback)
                    
                    # 更新缓冲区
                    await self._update_buffer(processed_chunk)
                
                # 控制发送速率
                await asyncio.sleep(0.01)
                
        except Exception as e:
            await self._emit_event(StreamEventType.ERROR, {"error": str(e)})
    
    async def _process_audio_chunk(self, audio_chunk: bytes) -> Optional[bytes]:
        """处理音频块"""
        try:
            if not audio_chunk:
                return None
            
            # 读取音频数据
            loop = asyncio.get_event_loop()
            audio_data, sample_rate = await loop.run_in_executor(
                self._executor,
                self._read_audio_data,
                audio_chunk
            )
            
            # 根据目标质量调整
            if sample_rate != self._current_quality.sample_rate:
                audio_data = await loop.run_in_executor(
                    self._executor,
                    self._resample_audio,
                    audio_data,
                    sample_rate,
                    self._current_quality.sample_rate
                )
            
            # 转换回字节
            processed_bytes = await loop.run_in_executor(
                self._executor,
                self._write_audio_data,
                audio_data,
                self._current_quality
            )
            
            return processed_bytes
            
        except Exception as e:
            logging.error(f"音频处理错误: {e}")
            return None
    
    def _read_audio_data(self, audio_bytes: bytes) -> tuple[np.ndarray, int]:
        """读取音频数据"""
        with io.BytesIO(audio_bytes) as buffer:
            data, sr = sf.read(buffer)
            return data, sr
    
    def _resample_audio(
        self,
        data: np.ndarray,
        orig_sr: int,
        target_sr: int
    ) -> np.ndarray:
        """重采样音频"""
        try:
            import librosa
            return librosa.resample(data, orig_sr=orig_sr, target_sr=target_sr)
        except ImportError:
            # 简单的线性插值重采样
            ratio = target_sr / orig_sr
            new_length = int(len(data) * ratio)
            indices = np.linspace(0, len(data) - 1, new_length)
            
            if len(data.shape) > 1:
                return np.array([
                    np.interp(indices, np.arange(len(data)), data[:, ch])
                    for ch in range(data.shape[1])
                ]).T
            else:
                return np.interp(indices, np.arange(len(data)), data)
    
    def _write_audio_data(self, data: np.ndarray, quality: StreamQuality) -> bytes:
        """写入音频数据"""
        with io.BytesIO() as buffer:
            sf.write(buffer, data, quality.sample_rate, format=quality.format)
            return buffer.getvalue()
    
    async def _send_to_client(
        self,
        audio_chunk: bytes,
        client_callback: Callable[[bytes], Any]
    ):
        """发送音频到客户端"""
        try:
            # 异步发送
            loop = asyncio.get_event_loop()
            await loop.run_in_executor(self._executor, client_callback, audio_chunk)
            
            # 更新统计
            self._stats["bytes_sent"] += len(audio_chunk)
            self._stats["chunks_sent"] += 1
            
        except Exception as e:
            logging.error(f"发送音频失败: {e}")
            await self._emit_event(StreamEventType.ERROR, {"error": str(e)})
    
    async def _update_buffer(self, audio_chunk: bytes):
        """更新缓冲区状态"""
        async with self._buffer_lock:
            self._buffer.append(audio_chunk)
            
            # 限制缓冲区大小
            if len(self._buffer) > self._max_buffer_size:
                self._buffer.pop(0)
            
            # 更新缓冲区健康度
            buffer_ratio = len(self._buffer) / self._max_buffer_size
            self._buffer_health = buffer_ratio
            
            # 发出缓冲区事件
            if buffer_ratio < 0.2:
                await self._emit_event(StreamEventType.BUFFER_LOW, {"buffer_ratio": buffer_ratio})
                self._stats["buffer_underruns"] += 1
            elif buffer_ratio > 0.8:
                await self._emit_event(StreamEventType.BUFFER_FULL, {"buffer_ratio": buffer_ratio})
            elif len(self._buffer) == 0:
                await self._emit_event(StreamEventType.BUFFER_EMPTY)
    
    async def _adaptive_control(self):
        """自适应质量控制"""
        while self._is_streaming:
            try:
                if self.config.enable_adaptive:
                    await self._adjust_quality()
                
                await asyncio.sleep(1.0)  # 每秒调整一次
                
            except Exception as e:
                logging.error(f"自适应控制错误: {e}")
    
    async def _adjust_quality(self):
        """调整流质量"""
        current_quality_index = self.config.quality_levels.index(self._current_quality)
        
        # 基于网络质量和缓冲区健康度调整
        combined_score = (self._network_quality * 0.6 + self._buffer_health * 0.4)
        
        if combined_score < 0.3:
            # 降低质量
            if current_quality_index > 0:
                new_index = current_quality_index - 1
                new_quality = self.config.quality_levels[new_index]
                await self._change_quality(new_quality, "network_or_buffer_low")
                
        elif combined_score > 0.8 and self._buffer_health > 0.7:
            # 提升质量
            if current_quality_index < len(self.config.quality_levels) - 1:
                new_index = current_quality_index + 1
                new_quality = self.config.quality_levels[new_index]
                await self._change_quality(new_quality, "network_and_buffer_good")
    
    async def _change_quality(self, new_quality: StreamQuality, reason: str):
        """改变流质量"""
        old_quality = self._current_quality
        self._current_quality = new_quality
        
        self._stats["quality_changes"] += 1
        
        # 调整缓冲区大小
        buffer_size_factor = new_quality.sample_rate / 16000  # 基准16kHz
        self._max_buffer_size = int(self.config.max_buffer_size * buffer_size_factor)
        
        await self._emit_event(StreamEventType.QUALITY_ADJUSTED, {
            "old_quality": old_quality.__dict__,
            "new_quality": new_quality.__dict__,
            "reason": reason
        })
    
    async def pause_streaming(self):
        """暂停流传输"""
        self._is_paused = True
    
    async def resume_streaming(self):
        """恢复流传输"""
        self._is_paused = False
    
    async def stop_streaming(self):
        """停止流传输"""
        self._is_streaming = False
        self._is_paused = False
    
    def add_event_handler(self, event_type: StreamEventType, handler: Callable):
        """添加事件处理器"""
        if event_type not in self._event_handlers:
            self._event_handlers[event_type] = []
        self._event_handlers[event_type].append(handler)
    
    def remove_event_handler(self, event_type: StreamEventType, handler: Callable):
        """移除事件处理器"""
        if event_type in self._event_handlers and handler in self._event_handlers[event_type]:
            self._event_handlers[event_type].remove(handler)
    
    async def _emit_event(self, event_type: StreamEventType, data: Dict[str, Any] = None):
        """发送事件"""
        event = StreamEvent(
            event_type=event_type,
            timestamp=time.time(),
            data=data or {}
        )
        
        await self._event_queue.put(event)
    
    async def _process_events(self):
        """处理事件"""
        while self._is_streaming:
            try:
                # 等待事件，超时1秒
                event = await asyncio.wait_for(self._event_queue.get(), timeout=1.0)
                
                # 调用事件处理器
                if event.event_type in self._event_handlers:
                    for handler in self._event_handlers[event.event_type]:
                        try:
                            if asyncio.iscoroutinefunction(handler):
                                await handler(event)
                            else:
                                handler(event)
                        except Exception as e:
                            logging.error(f"事件处理器错误: {e}")
                
            except asyncio.TimeoutError:
                # 超时是正常的，继续循环
                continue
            except Exception as e:
                logging.error(f"事件处理错误: {e}")
    
    def update_network_quality(self, quality: float):
        """更新网络质量 (0.0-1.0)"""
        self._network_quality = max(0.0, min(1.0, quality))
    
    def get_stats(self) -> Dict[str, Any]:
        """获取统计信息"""
        return {
            **self._stats,
            "current_quality": self._current_quality.__dict__,
            "network_quality": self._network_quality,
            "buffer_health": self._buffer_health,
            "buffer_size": len(self._buffer),
            "is_streaming": self._is_streaming,
            "is_paused": self._is_paused
        }
    
    async def cleanup(self):
        """清理资源"""
        self._is_streaming = False
        self._is_paused = False
        
        # 清空缓冲区
        async with self._buffer_lock:
            self._buffer.clear()
        
        # 清空事件队列
        while not self._event_queue.empty():
            try:
                self._event_queue.get_nowait()
            except asyncio.QueueEmpty:
                break
        
        # 关闭线程池
        self._executor.shutdown(wait=True)
        
        # 清空事件处理器
        self._event_handlers.clear()


class StreamQualityMonitor:
    """流质量监控器"""
    
    def __init__(self):
        self._latency_samples: List[float] = []
        self._bandwidth_samples: List[float] = []
        self._packet_loss_samples: List[float] = []
        self._max_samples = 100
    
    def record_latency(self, latency_ms: float):
        """记录延迟"""
        self._latency_samples.append(latency_ms)
        if len(self._latency_samples) > self._max_samples:
            self._latency_samples.pop(0)
    
    def record_bandwidth(self, bandwidth_kbps: float):
        """记录带宽"""
        self._bandwidth_samples.append(bandwidth_kbps)
        if len(self._bandwidth_samples) > self._max_samples:
            self._bandwidth_samples.pop(0)
    
    def record_packet_loss(self, loss_rate: float):
        """记录丢包率"""
        self._packet_loss_samples.append(loss_rate)
        if len(self._packet_loss_samples) > self._max_samples:
            self._packet_loss_samples.pop(0)
    
    def get_quality_score(self) -> float:
        """计算质量分数 (0.0-1.0)"""
        # 延迟分数 (越低越好)
        latency_score = 1.0
        if self._latency_samples:
            avg_latency = sum(self._latency_samples) / len(self._latency_samples)
            # 200ms以下为满分
            latency_score = max(0.0, 1.0 - (avg_latency - 50) / 150)
        
        # 带宽分数 (越高越好)
        bandwidth_score = 0.5  # 默认值
        if self._bandwidth_samples:
            avg_bandwidth = sum(self._bandwidth_samples) / len(self._bandwidth_samples)
            # 256kbps以上为满分
            bandwidth_score = min(1.0, avg_bandwidth / 256)
        
        # 丢包率分数 (越低越好)
        packet_loss_score = 1.0
        if self._packet_loss_samples:
            avg_loss = sum(self._packet_loss_samples) / len(self._packet_loss_samples)
            # 5%以下为满分
            packet_loss_score = max(0.0, 1.0 - avg_loss / 0.05)
        
        # 综合分数
        return (latency_score * 0.4 + bandwidth_score * 0.4 + packet_loss_score * 0.2)
    
    def get_metrics(self) -> Dict[str, Any]:
        """获取监控指标"""
        return {
            "latency_ms": {
                "current": self._latency_samples[-1] if self._latency_samples else 0,
                "average": sum(self._latency_samples) / len(self._latency_samples) if self._latency_samples else 0,
                "min": min(self._latency_samples) if self._latency_samples else 0,
                "max": max(self._latency_samples) if self._latency_samples else 0
            },
            "bandwidth_kbps": {
                "current": self._bandwidth_samples[-1] if self._bandwidth_samples else 0,
                "average": sum(self._bandwidth_samples) / len(self._bandwidth_samples) if self._bandwidth_samples else 0,
                "min": min(self._bandwidth_samples) if self._bandwidth_samples else 0,
                "max": max(self._bandwidth_samples) if self._bandwidth_samples else 0
            },
            "packet_loss_rate": {
                "current": self._packet_loss_samples[-1] if self._packet_loss_samples else 0,
                "average": sum(self._packet_loss_samples) / len(self._packet_loss_samples) if self._packet_loss_samples else 0,
                "min": min(self._packet_loss_samples) if self._packet_loss_samples else 0,
                "max": max(self._packet_loss_samples) if self._packet_loss_samples else 0
            },
            "quality_score": self.get_quality_score()
        }
